import akshare as ak
import numpy as np
import pandas as pd
from torch.utils import data

from StockDataset import StockDataset


def getDataset(dataset, lookback):
    X, y = [], []
    for i in range(len(dataset) - lookback):
        feature = dataset.iloc[i:i + lookback, :-1]
        target = dataset.iloc[i:i + lookback, -1:]
        X.append(feature)
        y.append(target)
    return X, y


class StockDataSource:

    def __init__(self, stock_code, days, ratio=0.8):
        self.stock_code = stock_code
        self.days = days
        self.stock_data_value = None
        self.ratio = ratio

    def getStockData(self):
        # get stock from akshare
        stock_his_data = ak.stock_zh_a_hist(symbol=self.stock_code, period="daily", start_date="20200101",
                                            adjust="qfq", timeout=50000)
        # sorted by date for desc
        stock_his_data = stock_his_data.sort_values("日期")
        # stock_data_source_data_frame
        stock_data = pd.DataFrame()
        stock_data["date"] = stock_his_data["日期"]
        stock_data["open"] = stock_his_data["开盘"]
        stock_data["close"] = stock_his_data["收盘"]
        stock_data["high"] = stock_his_data["最高"]
        stock_data["low"] = stock_his_data["最低"]
        stock_data["vol"] = stock_his_data["成交量"]
        stock_data["vol_amount"] = stock_his_data["成交额"]
        stock_data["change"] = stock_his_data["换手率"]
        stock_data["increment"] = stock_his_data["涨跌幅"]
        # add zz, zl1, zl2
        stock_data_new = pd.DataFrame(stock_data)
        stock_data_new["zz"] = (stock_data_new["open"] + stock_data_new["low"] + stock_data_new["high"]) / 3
        stock_data_new["zl1"] = stock_data_new["zz"].ewm(span=7, adjust=True).mean().round(2)
        stock_data_new["zl2"] = stock_data_new["zz"].ewm(span=55, adjust=True).mean().round(2)
        stock_data_new = stock_data_new.drop('zz', axis=1)
        # ref 5 price (include vol, change)
        stock_data_new["vol1"] = (stock_data_new["vol"] - (stock_data_new["vol"] +
                                                           stock_data_new["vol"].shift(1) +
                                                           stock_data_new["vol"].shift(2) +
                                                           stock_data_new["vol"].shift(3) +
                                                           stock_data_new["vol"].shift(4)) / 5) / stock_data_new["vol"]

        stock_data_new["change1"] = (stock_data_new["change"] - (stock_data_new["change"] +
                                                                 stock_data_new["change"].shift(1) +
                                                                 stock_data_new["change"].shift(2) +
                                                                 stock_data_new["change"].shift(3) +
                                                                 stock_data_new["change"].shift(4)) / 5) / \
                                    stock_data_new["change"]

        # add stock flag value
        stock_data_new["flag"] = np.zeros(len(stock_data_new))
        # set flag value
        stock_data_new = self.setFlagValue(stock_data_new)
        # drop column
        stock_data_new = stock_data_new.drop('vol_amount', axis=1)
        stock_data_new = stock_data_new.drop('date', axis=1)
        # setting stock data value
        self.stock_data_value = stock_data_new

        stock_data_new["open"] = (stock_data_new["open"] - stock_data_new["zl1"]) / stock_data_new["open"]
        stock_data_new["close"] = (stock_data_new["close"] - stock_data_new["zl1"]) / stock_data_new["close"]
        stock_data_new["high"] = (stock_data_new["high"] - stock_data_new["zl1"]) / stock_data_new["high"]
        stock_data_new["low"] = (stock_data_new["low"] - stock_data_new["zl1"]) / stock_data_new["low"]

        self.stock_data_value = stock_data_new[["open", "close", "high", "low", "vol1", "change1", "flag"]]

        return self.stock_data_value

    def setFlagValue(self, stock_data):
        inc_percent5 = 0.05
        inc_percent4 = 0.04
        inc_percent3 = 0.03
        inc_percent2 = 0.02
        inc_percent1 = 0.01
        for index in range(0, len(stock_data)):
            next_value = stock_data.loc[index + 1:index + self.days]["close"].max()
            current_close_value = stock_data.loc[index]["close"]
            if (next_value - current_close_value) / current_close_value > inc_percent5:
                stock_data.iloc[index, 13] = 0.5
            elif (next_value - current_close_value) / current_close_value > inc_percent4:
                stock_data.iloc[index, 13] = 0.4
            elif (next_value - current_close_value) / current_close_value > inc_percent3:
                stock_data.iloc[index, 13] = 0.3
            elif (next_value - current_close_value) / current_close_value > inc_percent2:
                stock_data.iloc[index, 13] = 0.2
            elif (next_value - current_close_value) / current_close_value > inc_percent1:
                stock_data.iloc[index, 13] = 0.1
            elif next_value >= current_close_value:
                stock_data.iloc[index, 13] = 0.05
            else:
                if (current_close_value - next_value) / current_close_value > inc_percent5:
                    stock_data.iloc[index, 13] = -0.5
                elif (current_close_value - next_value) / current_close_value > inc_percent4:
                    stock_data.iloc[index, 13] = -0.4
                elif (current_close_value - next_value) / current_close_value > inc_percent3:
                    stock_data.iloc[index, 13] = -0.3
                elif (current_close_value - next_value) / current_close_value > inc_percent2:
                    stock_data.iloc[index, 13] = -0.2
                elif (current_close_value - next_value) / current_close_value > inc_percent1:
                    stock_data.iloc[index, 13] = -0.1
                elif current_close_value > next_value:
                    stock_data.iloc[index, 13] = -0.05
        return stock_data

    def getTrainData(self):
        if self.stock_data_value is None:
            self.getStockData()

        if self.stock_data_value is not None:
            train_len = int(len(self.stock_data_value) * self.ratio)
            test_len = len(self.stock_data_value) - train_len
            train_data = self.stock_data_value[:-test_len]
            test_data = self.stock_data_value[-test_len:]

            train_data_loader1 = StockDataset(train_data, 5)
            test_data_loader1 = StockDataset(test_data, 5)

            train_data_loader = data.DataLoader(train_data_loader1, batch_size=5, shuffle=False, drop_last=False)
            test_data_loader = data.DataLoader(test_data_loader1, batch_size=5, shuffle=False, drop_last=False)
            return train_data_loader, test_data_loader
        return None
